package com.orientechnologies.orient.distributed.context;

import com.orientechnologies.orient.core.transaction.ONodeId;
import com.orientechnologies.orient.core.transaction.OTransactionIdPromise;
import com.orientechnologies.orient.distributed.context.coordination.result.OAcceptResult;
import com.orientechnologies.orient.distributed.context.coordination.result.OQuorumNotReached;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Optional;
import java.util.Set;

public class OResponseCollector {

  public record CompleteInfo(
      OCompleteAction action,
      OTransactionIdPromise promise,
      Set<ONodeId> nodes,
      Optional<OAcceptResult> result) {}
  ;

  private final OCompleteAction action;
  private final int quorum;
  private final OTransactionIdPromise promise;
  private final Set<ONodeId> expected;
  private final Set<ONodeId> toReceive;
  private final Set<ONodeId> success = new HashSet<>();
  private final Set<ONodeId> lost = new HashSet<>();
  private final Set<ONodeId> failure = new HashSet<>();
  private final Map<OAcceptResult, Integer> results = new HashMap<>();

  public OResponseCollector(
      OCompleteAction action, OTransactionIdPromise promise, int quorum, Set<ONodeId> activeNodes) {
    this.action = action;
    this.promise = promise;
    this.quorum = quorum;
    this.expected = activeNodes;
    this.toReceive = new HashSet<>(activeNodes);
  }

  public Optional<CompleteInfo> receive(ONodeId node) {
    if (toReceive.remove(node)) {
      success.add(node);
      if (success.size() == quorum) {
        return Optional.of(new CompleteInfo(action, promise, expected, Optional.empty()));
      }
    }
    return Optional.empty();
  }

  public boolean hasSuccessQuorum() {
    return success.size() >= quorum;
  }

  public boolean hasFailureQuorum() {
    return lost.size() + failure.size() >= quorum;
  }

  public boolean isFinished() {
    return expected.size() == success.size() + lost.size() + failure.size();
  }

  public boolean isFinishedNotQuorum() {
    return isFinished() && !hasFailureQuorum() && !hasSuccessQuorum();
  }

  public Set<ONodeId> getExpected() {
    return expected;
  }

  public Optional<CompleteInfo> disconnected(ONodeId node) {
    if (toReceive.remove(node)) {
      lost.add(node);
      if (lost.size() + failure.size() == quorum) {
        OAcceptResult result = computeResult();
        return Optional.of(new CompleteInfo(action, promise, expected, Optional.of(result)));
      }
      if (isFinishedNotQuorum()) {
        OAcceptResult result = computeResult();
        return Optional.of(new CompleteInfo(action, promise, expected, Optional.of(result)));
      }
    }
    return Optional.empty();
  }

  public Optional<CompleteInfo> fail(ONodeId node, OAcceptResult acceptResult) {
    if (toReceive.remove(node)) {
      failure.add(node);
      this.results.compute(
          acceptResult,
          (k, v) -> {
            if (v == null) {
              return 1;
            } else {
              return v + 1;
            }
          });
      if (lost.size() + failure.size() == quorum) {
        OAcceptResult result = computeResult();
        return Optional.of(new CompleteInfo(action, promise, expected, Optional.of(result)));
      }
      if (isFinishedNotQuorum()) {
        OAcceptResult result = computeResult();
        return Optional.of(new CompleteInfo(action, promise, expected, Optional.of(result)));
      }
    }
    return Optional.empty();
  }

  private OAcceptResult computeResult() {
    for (var res : this.results.entrySet()) {
      if (res.getValue() >= quorum) {
        return res.getKey();
      }
    }
    return new OQuorumNotReached(this.results.keySet());
  }
}
